#include "memory.h"
#include "uart.h"

// 内存池起始地址和大小
#define MEMORY_START 0x80100000
#define MEMORY_SIZE  0x10000    // 64KB

// 最小分配单位
#define MIN_ALLOC_SIZE 16

static struct mem_block *head = NULL;

// 用于数字转字符串的辅助函数
static void print_num(size_t num) {
    char buf[32];
    int i = 0;
    
    if(num == 0) {
        uart_putc('0');
        return;
    }
    
    while(num > 0) {
        buf[i++] = '0' + (num % 10);
        num /= 10;
    }
    
    while(--i >= 0) {
        uart_putc(buf[i]);
    }
}

void memory_init(void) {
    head = (struct mem_block*)MEMORY_START;
    head->size = MEMORY_SIZE - sizeof(struct mem_block);
    head->is_free = 1;
    head->next = NULL;
    
    uart_puts("Memory management initialized.\n");
    uart_puts("Total memory: ");
    print_num(MEMORY_SIZE);
    uart_puts(" bytes\n");
}

void* memory_alloc(size_t size) {
    struct mem_block *current = head;
    struct mem_block *best_fit = NULL;
    size_t min_size_diff = MEMORY_SIZE;
    
    // 对齐到最小分配单位
    size = (size + MIN_ALLOC_SIZE - 1) & ~(MIN_ALLOC_SIZE - 1);
    
    // 查找最佳适配块
    while(current != NULL) {
        if(current->is_free && current->size >= size) {
            size_t size_diff = current->size - size;
            if(size_diff < min_size_diff) {
                min_size_diff = size_diff;
                best_fit = current;
            }
        }
        current = current->next;
    }
    
    if(best_fit == NULL) {
        return NULL;  // 没有足够的内存
    }
    
    // 如果剩余空间足够大，分割块
    if(min_size_diff > sizeof(struct mem_block) + MIN_ALLOC_SIZE) {
        struct mem_block *new_block = (struct mem_block*)((char*)best_fit + sizeof(struct mem_block) + size);
        new_block->size = best_fit->size - size - sizeof(struct mem_block);
        new_block->is_free = 1;
        new_block->next = best_fit->next;
        
        best_fit->size = size;
        best_fit->next = new_block;
    }
    
    best_fit->is_free = 0;
    return (void*)((char*)best_fit + sizeof(struct mem_block));
}

void memory_free(void* ptr) {
    if(!ptr) return;
    
    struct mem_block *block = (struct mem_block*)((char*)ptr - sizeof(struct mem_block));
    block->is_free = 1;
    
    // 合并相邻的空闲块
    struct mem_block *current = head;
    while(current != NULL) {
        if(current->is_free && current->next && current->next->is_free) {
            current->size += sizeof(struct mem_block) + current->next->size;
            current->next = current->next->next;
            continue;
        }
        current = current->next;
    }
}

void memory_stats(void) {
    struct mem_block *current = head;
    size_t total_free = 0;
    size_t total_used = 0;
    int free_blocks = 0;
    int used_blocks = 0;
    
    uart_puts("\nMemory Statistics:\n");
    uart_puts("================\n");
    
    while(current != NULL) {
        if(current->is_free) {
            total_free += current->size;
            free_blocks++;
        } else {
            total_used += current->size;
            used_blocks++;
        }
        current = current->next;
    }
    
    uart_puts("Free memory: ");
    print_num(total_free);
    uart_puts(" bytes in ");
    print_num(free_blocks);
    uart_puts(" blocks\n");
    
    uart_puts("Used memory: ");
    print_num(total_used);
    uart_puts(" bytes in ");
    print_num(used_blocks);
    uart_puts(" blocks\n");
}

void memory_test(void) {
    uart_puts("\nRunning memory test...\n");
    
    // 测试1：分配和释放
    uart_puts("Test 1: Allocating memory blocks...\n");
    void *ptr1 = memory_alloc(100);
    uart_puts("- Allocated 100 bytes at: 0x");
    uart_put_hex((unsigned long)ptr1);
    uart_puts("\n");
    
    void *ptr2 = memory_alloc(200);
    uart_puts("- Allocated 200 bytes at: 0x");
    uart_put_hex((unsigned long)ptr2);
    uart_puts("\n");
    
    uart_puts("\nTest 2: Freeing memory blocks...\n");
    uart_puts("- Freeing first allocation...\n");
    memory_free(ptr1);
    
    uart_puts("- Freeing second allocation...\n");
    memory_free(ptr2);
    
    // 测试2：尝试分配大块内存
    uart_puts("\nTest 3: Attempting large allocation...\n");
    void *ptr3 = memory_alloc(MEMORY_SIZE);
    if(ptr3 == NULL) {
        uart_puts("- Large allocation failed as expected\n");
    }
    
    uart_puts("\nMemory test completed successfully!\n");
}
